#!/usr/bin/env ./venv/bin/python
# -*- coding: utf-8 -*-
import os
import cv2
import numpy as np
import torch
import threading
import queue
import time
import math
import random
import asyncio
import websockets
import json
from ultralytics import YOLO          # YOLOv8 官方库
from PIL import Image

# ---------------- 参数 ----------------
RTSP_URL = "rtsp://10.100.179.102:8554/d435i"
PORT       = 9002
FREQ_HZ    = 50          # 飞控指令频率
THRESHOLD  = 0.7          # 置信度阈值
NMS_THRES  = 0.45
DEVICE     = "cuda" if torch.cuda.is_available() else "cpu"

# 线程间通信
frame_que   = queue.Queue(maxsize=2)   # 最新 2 帧，满就丢
result_que  = queue.Queue(maxsize=2)   # 画框后的帧
latest      = {"x": 0.0, "y": 0.0, "z": 0.0, "yaw": 0.0}
lock        = threading.Lock()
shutdown    = threading.Event()        # 线程安全的事件对象

# ---------------- 1. 加载离线模型 ----------------
weight_path = "./models/yolo/yolov8n.pt"
if not os.path.exists(weight_path):
    raise FileNotFoundError("请先下载 yolov8n.pt 放到 ./models/ 目录")
model = YOLO(weight_path)
model.to(DEVICE)

# 读取类别名
with open("./models/coco.names") as f:
    CLASS_NAMES = [line.strip() for line in f]
PERSON_ID = CLASS_NAMES.index("person")

# ---------------- 2. 推理函数 ----------------
def yolo_infer(frame):
    """
    返回：画好框的 frame
    """
    # Ultralytics 支持直接传 numpy；返回 Results 列表
    results = model.predict(
        source=frame,
        conf=THRESHOLD,
        iou=NMS_THRES,
        device=DEVICE,
        verbose=False,
        classes=[PERSON_ID]      # 只检人
    )[0]

    # 画框
    annotated = results.plot()               
    return annotated

# ---------------- 3. 读流线程 ----------------
def capture_task():
    cap = cv2.VideoCapture(RTSP_URL + "?fflags=discardcorrupt", cv2.CAP_FFMPEG)
    cap.set(cv2.CAP_PROP_BUFFERSIZE, 1)

    while not shutdown.is_set():      # ← 检测标志
        ret, frame = cap.read()
        if not ret:
            time.sleep(0.5)
            continue
        if frame_que.full():
            frame_que.get_nowait()
        frame_que.put(frame)

    # ----- 清理 -----
    cap.release()

# ---------------- 4. 推理线程 ----------------
def inference_task():
    cnt = 0
    while not shutdown.is_set():
        try:
            frame = frame_que.get()
        except queue.Empty:
            continue

        # cnt += 1
        # if cnt % 3 != 0:        # 抽帧，约 10 Hz
        #     continue

        out_frame = yolo_infer(frame)

        with lock:
            latest.update(zip(("x","y","z","yaw"), get_cmd(frame)))

        if result_que.full():
            result_que.get_nowait()
        result_que.put(out_frame)
        # print(f"infer cost:{(time.time()-t0)*1000:.1f} ms")  # about 35ms

# ---------------- 5. 显示线程 ----------------
def display_task():
    while not shutdown.is_set():
        try:
            frame = result_que.get()
        except queue.Empty:
            continue

        cv2.imshow("D435i-YOLO", frame)
        if cv2.waitKey(1) & 0xFF == ord('q'):
            shutdown.set()        # 触发退出
            break

    cv2.destroyAllWindows()

# ---------------- 6. WebSocket ----------------
async def ws_handler(websocket, path):
    print("client connected")
    try:
        while True:
            with lock:
                payload = json.dumps(latest)
            await websocket.send(payload)
            await asyncio.sleep(1 / FREQ_HZ)
    except websockets.exceptions.ConnectionClosed:
        print("client disconnected")

async def ws_main():
    stop = asyncio.Future()
    # shutdown.add_done_callback(lambda _: stop.set_result(None))  # 联动
    async with websockets.serve(ws_handler, "0.0.0.0", PORT):
        print(f"WS server listening on :{PORT}")
        await stop          # 收到全局退出信号就结束

# ---------------- 7. 占位算法 ----------------
def get_cmd(image):
    """替换成你的视觉算法"""
    return (random.uniform(-10, 10),
            random.uniform(-10, 10),
            random.uniform(0, 5),
            random.uniform(-math.pi, math.pi))

# ---------------- 8. 主入口 ----------------
if __name__ == "__main__":
    th_cap = threading.Thread(target=capture_task, daemon=True)
    th_inf = threading.Thread(target=inference_task, daemon=True)
    th_disp = threading.Thread(target=display_task, daemon=True)
    th_cap.start()
    th_inf.start()
    th_disp.start()

    try:
        # 把 WebSocket 也跑在同一个线程池
        asyncio.run(ws_main())
    except KeyboardInterrupt:
        print("\nCtrl-C caught, shutting down...")
    finally:
        shutdown.set()          # 通知所有线程
        th_cap.join(timeout=1)
        th_inf.join(timeout=1)
        th_disp.join(timeout=1)